package com.huakai.springenv.perf;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

public class TrafficRoutingDataSource extends AbstractRoutingDataSource {

    // ThreadLocal 用于存储当前请求是否是压测流量
    private static final ThreadLocal<Boolean> PRESSURE_TEST_FLAG = new ThreadLocal<>();

    // 设置压测标识
    public static void setPressureTestFlag(boolean isPressureTest) {
        PRESSURE_TEST_FLAG.set(isPressureTest);
    }

    // 清除压测标识
    public static void clearPressureTestFlag() {
        PRESSURE_TEST_FLAG.remove();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        // 根据 ThreadLocal 判断是否是压测流量
        Boolean isPressureTest = PRESSURE_TEST_FLAG.get();
        if (Boolean.TRUE.equals(isPressureTest)) {
            return "shadowDataSource";  // 返回影子库的数据源 key
        } else {
            return "primaryDataSource";  // 返回原始库的数据源 key
        }
    }
}
